#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
QTorch量化交易框架 - 统一API演示
支持自动GPU加速的量化交易回测
"""

import logging
import argparse
import os
import yaml
import pandas as pd
import time
from datetime import datetime
from qtorch.core.qengine import QEngine
from qtorch.utils.device_utils import get_device, is_gpu_available

# 导入策略
from qtorch.strategy.dual_moving_average import DualMovingAverageStrategy
from qtorch.strategy.simple_sma_strategy import SimpleSMAStrategy
from qtorch.strategy.simple_macd_strategy import SimpleMACDStrategy

# 如果需要其他原有策略，导入它们
try:
    from qtorch.strategy.macd_strategy import MACDStrategy
    from qtorch.strategy.rsi_strategy import RSIStrategy
    from qtorch.strategy.bollinger_bands import BollingerBandsStrategy
except ImportError:
    # 如果有些策略没有实现，只是记录警告
    logging.warning("部分策略模块未找到，将使用可用的策略")

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='QTorch量化交易回测系统 - 支持GPU加速')
    parser.add_argument('--config', type=str, help='配置文件路径')
    parser.add_argument('--strategy', type=str, default='dual_ma', help='策略名称')
    parser.add_argument('--symbol', type=str, default='000001', help='交易标的代码')
    parser.add_argument('--start', type=str, default='2020-01-01', help='回测开始日期')
    parser.add_argument('--end', type=str, default='2022-12-31', help='回测结束日期')
    parser.add_argument('--params', type=str, help='策略参数，格式：param1=value1,param2=value2')
    parser.add_argument('--output', type=str, default='./output', help='结果输出路径')
    parser.add_argument('--cash', type=float, default=1000000, help='初始资金')
    parser.add_argument('--commission', type=float, default=0.0003, help='交易佣金')
    
    # GPU加速相关参数
    parser.add_argument('--force-cpu', action='store_true', help='强制使用CPU，即使有GPU可用')
    parser.add_argument('--gpu-stats', action='store_true', help='显示GPU性能统计')
    parser.add_argument('--mixed-precision', action='store_true', help='使用混合精度计算')
    parser.add_argument('--benchmark', action='store_true', help='运行性能基准测试')
    
    return parser.parse_args()

def parse_strategy_params(params_str):
    """解析策略参数字符串"""
    if not params_str:
        return {}
        
    param_dict = {}
    for param in params_str.split(','):
        key, value = param.split('=')
        param_dict[key] = float(value) if '.' in value else int(value)
    return param_dict

def get_strategy_class(strategy_name):
    """获取策略类"""
    strategy_map = {
        'dual_ma': DualMovingAverageStrategy,
        'simple_sma': SimpleSMAStrategy,
        'simple_macd': SimpleMACDStrategy,
        # 兼容旧版策略名称
        'macd': MACDStrategy if 'MACDStrategy' in globals() else SimpleMACDStrategy,
        'rsi': RSIStrategy if 'RSIStrategy' in globals() else None,
        'bollinger': BollingerBandsStrategy if 'BollingerBandsStrategy' in globals() else None
    }
    
    if strategy_name not in strategy_map:
        logging.warning(f"未找到策略 '{strategy_name}'，使用默认策略 'dual_ma'")
        return strategy_map['dual_ma']
        
    if strategy_map[strategy_name] is None:
        logging.warning(f"策略 '{strategy_name}' 未实现，使用默认策略 'dual_ma'")
        return strategy_map['dual_ma']
        
    return strategy_map[strategy_name]

def show_gpu_info():
    """显示GPU信息"""
    if is_gpu_available():
        device = get_device()
        logging.info(f"已检测到可用GPU: {device}")
        
        # 如果是CUDA设备，显示更多信息
        if hasattr(device, 'type') and device.type == 'cuda':
            import torch
            for i in range(torch.cuda.device_count()):
                props = torch.cuda.get_device_properties(i)
                logging.info(f"  CUDA:{i} - {props.name}, 内存: {props.total_memory / 1024**3:.1f}GB")
    else:
        logging.info("未检测到GPU，将使用CPU进行计算")

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_args()
    
    # 显示GPU信息
    show_gpu_info()
    
    start_time = time.time()
    
    try:
        # 创建回测引擎实例
        engine = QEngine(force_cpu=args.force_cpu)
        
        # 设置GPU相关参数
        if not args.force_cpu and is_gpu_available():
            engine.set_mixed_precision(args.mixed_precision)
            if args.benchmark:
                engine.enable_benchmark_mode(True)
                logging.info("启用性能对比模式")
        
        # 如果提供了配置文件，使用配置文件
        if args.config:
            engine.set_config(args.config)
        
        # 添加交易数据
        symbols = args.symbol.split(',')
        for symbol in symbols:
            engine.add_data(
                symbol=symbol,
                start_date=args.start,
                end_date=args.end
            )
            
        # 设置回测参数
        engine.set_cash(args.cash)
        engine.set_commission(args.commission)
        
        # 解析策略参数
        strategy_params = parse_strategy_params(args.params)
        
        # 获取策略类
        strategy_class = get_strategy_class(args.strategy)
        
        # 添加策略
        engine.add_strategy(strategy_class, **strategy_params)
        
        # 运行回测
        logging.info(f"开始回测: 策略={args.strategy}, 股票={args.symbol}, 时间={args.start}至{args.end}")
        result = engine.run()
        
        # 打印回测结果摘要
        result.summary()
        
        # 如果启用了GPU统计
        if args.gpu_stats:
            result.print_performance_stats()
        
        # 计算总运行时间
        total_time = time.time() - start_time
        logging.info(f"回测总耗时: {total_time:.2f}秒")
        
        # 生成回测报告
        output_dir = args.output
        os.makedirs(output_dir, exist_ok=True)
        result.plot(output_dir=output_dir, show_gpu_stats=args.gpu_stats)
        
        # 保存回测数据和指标
        metrics_path = os.path.join(output_dir, "metrics.yaml")
        with open(metrics_path, 'w', encoding='utf-8') as f:
            yaml.dump(result.metrics, f, default_flow_style=False, allow_unicode=True)
            
        result_path = os.path.join(output_dir, "result.csv")
        result.data.to_csv(result_path)
        
        logging.info(f"结果已保存到：{output_dir}")
        logging.info(f"回测报告已生成：{os.path.join(output_dir, 'report.html')}")
        
    except Exception as e:
        logging.error(f"运行回测失败: {str(e)}", exc_info=True)
        raise

if __name__ == '__main__':
    main()